import random
import math
import copy

def assign_cluster(x, centroids):
    min_dist = float('inf')
    cluster_id = 0
    for i, c in enumerate(centroids):
        # 合并距离计算，保持开根号
        dist = math.sqrt(sum((x[j] - c[j]) ** 2 for j in range(len(x))))
        if dist < min_dist:
            min_dist = dist
            cluster_id = i
    return cluster_id

def Kmeans(data, k, epsilon=1e-4, iteration=100):
    if not data or k <= 0 or k > len(data):
        raise ValueError("数据为空、k 无效或 k 大于样本数")

    dimension = len(data[0])
    centroids = random.sample(data, k)

    for iter_count in range(iteration):
        # 分配簇
        clusters = [assign_cluster(point, centroids) for point in data]
        # 保存旧中心
        old_centroids = copy.deepcopy(centroids)
        # 用列表推导式简化簇中心更新，保持逻辑清晰
        for i in range(k):
            cluster_points = [data[j] for j, lbl in enumerate(clusters) if lbl == i]
            if cluster_points:
                centroids[i] = [sum(p[dim] for p in cluster_points) / len(cluster_points)
                                for dim in range(dimension)]
            else:
                centroids[i] = random.choice(data)

        # 简化收敛判断的距离计算
        total_move = sum(math.sqrt(sum((oc[d] - nc[d]) ** 2 for d in range(dimension)))
                        for oc, nc in zip(old_centroids, centroids))

        if total_move < epsilon:
            print(f"第 {iter_count + 1} 次迭代后收敛")
            break
    else:
        print(f"达到最大迭代次数 {iteration}")

    return centroids, clusters

if __name__ == "__main__":
    data = [
        [1, 2], [1, 4], [1, 0], [10, 2], [10, 4],
        [10, 0], [5, 8], [6, 8], [7, 8], [2, 2],
        [1, 3], [0, 1], [11, 3], [9, 1], [10, 5]
    ]

    print("开始 K-means 聚类测试...")
    print(f"数据点数量: {len(data)} | 特征维度: {len(data[0])}")
    print("-" * 50)

    k = 3
    final_centroids, cluster_labels = Kmeans(data, k=k)

    # 简化结果输出格式，保留关键信息
    print("\n最终簇中心:")
    for i, center in enumerate(final_centroids):
        print(f"  簇 {i}: [{center[0]:.3f}, {center[1]:.3f}]")

    print("\n各簇包含的点:")
    for i in range(k):
        points = [str(data[j]) for j, lbl in enumerate(cluster_labels) if lbl == i]
        print(f"  簇 {i} ({len(points)} 个): {', '.join(points)}")